﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Threading;

namespace DBRandomGenerator
{
    public class ThreadPoolWait : IDisposable
    {
       // Notice that the counter's value is initially set to one rather than to zero, which may seem odd. 
        //This is done to work around the problem of a work item finishing before all of the work items have 
        //been queued; in such a scenario, the event could be signaled before all work items have completed.
        //As a workaround, the counter is initially set to one to ensure that the counter remains above 
        //zero until WaitOne is called. WaitOne begins by removing this extra value so that the count 
        //accurately represents the remaining work items to be completed.
        private int _remainingWorkItems = 1;
        private ManualResetEvent _done = new ManualResetEvent(false);

        public void QueueUserWorkItem(WaitCallback callback)
        {
            QueueUserWorkItem(callback, null);
        }

        public void QueueUserWorkItem(WaitCallback callback, object state)
        {
            ThrowIfDisposed();
            QueuedCallback qc = new QueuedCallback();
            qc.Callback = callback;
            qc.State = state;
            lock (_done) _remainingWorkItems++;
            ThreadPool.QueueUserWorkItem(new WaitCallback(HandleWorkItem), qc);
        }

        //Note, too, that by the time WaitOne is called, all of the work items may have completed and thus
        // the call WaitOne makes to decrement the counter may, in fact, bring the counter to zero. Thus,
        //WaitOne has to perform the same check as the work item, signaling the ManualResetEvent if there
        //is no outstanding work to be done.
        public bool WaitOne() { return WaitOne(-1, false); }

        public bool WaitOne(TimeSpan timeout, bool exitContext)
        {
            return WaitOne((int) timeout.TotalMilliseconds, exitContext);
        }

        public bool WaitOne(int millisecondsTimeout, bool exitContext)
        {
            ThrowIfDisposed();
            DoneWorkItem();
            bool rv = _done.WaitOne(millisecondsTimeout, exitContext);
            lock (_done)
            {
                if (rv)
                {
                    _remainingWorkItems = 1;
                    _done.Reset();
                }
                else _remainingWorkItems++;
            }
            return rv;
        }

        private void HandleWorkItem(object state)
        {
            QueuedCallback qc = (QueuedCallback) state;
            try { qc.Callback(qc.State); }
            finally { DoneWorkItem(); }
        }

        private void DoneWorkItem()
        {
            lock (_done)
            {
                --_remainingWorkItems;
                if (_remainingWorkItems == 0) _done.Set();
            }
        }

        private class QueuedCallback
        {
            public WaitCallback Callback;
            public object State;
        }

        private void ThrowIfDisposed()
        {
            if (_done == null) throw new ObjectDisposedException(GetType().Name);
        }

        public void Dispose()
        {
            if (_done != null)
            {
                ((IDisposable) _done).Dispose();
                _done = null;
            }
        }
    } 

}
